package com.noahsloan.nutils.bytes;

import java.io.Closeable;
import java.io.IOException;
import java.io.OutputStream;
import java.util.BitSet;

/**
 * Allows writing to an OutputStream one bit at a time.
 * 
 * Buffers 1 byte at a time.
 * 
 * @author noah
 * 
 */
public class BitWriter implements Closeable {

    /**
     * Essentially the max index of the buffer, which is a single byte, so 7 (in
     * general).
     */
    public static final int MAX_FILL = Byte.SIZE - 1;

    private OutputStream out; // stream to write bits to

    private int buffer; // the byte being buffered

    private int bufferFill; // how much of the buffer has been used

    /**
     * Creates a BitWriter wrapped around out.
     * 
     * @param out
     */
    public BitWriter(OutputStream out) {
        this.out = out;
        buffer = 0;
        bufferFill = 0;
    }

    /**
     * Write a bit.
     * 
     * @param bit
     * @throws IOException
     */
    public void write(boolean bit) throws IOException {
        if (bit) {
            buffer |= 1 << MAX_FILL - bufferFill;
        }
        bufferFill++;
        if (bufferFill > MAX_FILL) {
            out.write(buffer);
            bufferFill = 0;
            buffer = 0;
        }
    }

    /**
     * Write a bunch of bits at once.
     * 
     * @param bits
     * @throws IOException
     */
    public void write(boolean[] bits) throws IOException {
        for (boolean bit : bits) {
            write(bit);
        }
    }

    /**
     * Writes the given byte to the stream.
     * 
     * @param b
     * @throws IOException
     */
    public void write(byte b) throws IOException {
        write(b, Byte.SIZE);
    }

    /**
     * Write the given bytes to the stream.
     * 
     * @param bytes
     * @throws IOException
     */
    public void write(byte[] bytes) throws IOException {
        for (byte b : bytes) {
            write(b);
        }
    }

    /**
     * Writes the given short to the stream.
     * 
     * @param s
     * @throws IOException
     */
    public void write(short s) throws IOException {
        write(s, Short.SIZE);
    }

    /**
     * Writes the given int to the stream.
     * 
     * @param i
     * @throws IOException
     */
    public void write(int i) throws IOException {
        write(i, Integer.SIZE);
    }

    /**
     * Writes the given float the the stream.
     * 
     * @param f
     * @throws IOException
     */
    public void write(float f) throws IOException {
        write(Float.floatToIntBits(f));
    }

    /**
     * Writes the given long to the stream.
     * 
     * @param l
     * @throws IOException
     */
    public void write(long l) throws IOException {
        write(l, Long.SIZE);
    }

    /**
     * Writes the given double to the stream.
     * 
     * @param d
     * @throws IOException
     */
    public void write(double d) throws IOException {
        this.write(Double.doubleToLongBits(d));
    }

    /**
     * Writes out the given bit set.
     * 
     * @param bits
     * @throws IOException
     */
    public void write(BitSet bits) throws IOException {
        for (int i = 0; i < bits.length(); i++) {
            write(bits.get(i));
        }
    }

    /**
     * Internal multi-bit writer. Writes the last size bits from l, highest
     * order bits first.
     * 
     * For example, a binary valued long of 10101100 (leading 0s removed) passed
     * with a size of 4 would result in 1100 being written.
     * 
     * @param l
     *            the value buffer
     * @param size
     * @throws IOException
     */
    protected void write(long l, int size) throws IOException {
        long bit = 1l << (size - 1);
        for (int j = 0; j < size; j++) {
            boolean on = (bit & l) != 0;
            write(on);
            bit = bit >>> 1;
        }
    }

    /**
     * Flushes and closes the underlying OutputStream.
     * 
     * @throws IOException
     *             if there is a problem writing to, flushing, or closing the
     *             OutputStream
     */
    public void close() throws IOException {
        flush();
        out.flush();
        out.close();
    }

    /**
     * Pads out the current byte with 0s.
     * 
     * @throws IOException
     */
    public void flush() throws IOException {
        while (bufferFill != 0) {
            write(false);
        }
    }
}
